package org.lib.postgresqlhello.record.shard;

import java.sql.Connection;

import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.reflection.DefaultReflectorFactory;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;

/**
 * 分表拦截器
 * @author ben
 * @date 2021-11-18 22:45:44 CST
 */
@Intercepts({
	@Signature(type = StatementHandler.class,
			method = "prepare",
			args = {Connection.class, Integer.class}
	)
})
public class TableShardInterceptor implements Interceptor {

	@Override
	public Object intercept(Invocation invocation) throws Throwable {
		StatementHandler sh = (StatementHandler) invocation.getTarget();
		MetaObject mo = MetaObject.forObject(sh, SystemMetaObject.DEFAULT_OBJECT_FACTORY, 
				SystemMetaObject.DEFAULT_OBJECT_WRAPPER_FACTORY, 
				new DefaultReflectorFactory());
		
		MappedStatement ms = (MappedStatement) mo.getValue("delegate.mappedStatement");
		
		String id = ms.getId();
		id = id.substring(0, id.lastIndexOf('.'));
		Class<?> cl = Class.forName(id);
		
		// 获取TableShard注释
		TableShard tableShard = cl.getAnnotation(TableShard.class);
		if (tableShard != null) {
			// 做分表
			String tableName = tableShard.tableName();
			Class<? extends ITableShardStrategy> strategyClazz = tableShard.shardStrategy();
			ITableShardStrategy strategy = strategyClazz.newInstance();
			String newTableName = strategy.tableShard(tableName);
			
			// 获取源sql
			String sql = (String) mo.getValue("delegate.boundSql.sql");
			// 用新sql代替旧sql
			mo.setValue("delegate.boundSql.sql", sql.replaceAll(tableName, newTableName));
		}
		
		// 传递给下一个拦截器
		return invocation.proceed();
	}

	@Override
	public Object plugin(Object target) {
		// 当目标类是StatementHandler类型时，才包装目标类，
		// 否者直接返回目标本身, 减少目标被代理的次数
		if (target instanceof StatementHandler) {
			return Plugin.wrap(target, this);
		} else {
			return target;
		}
	}
  
}
